{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QUANWN3rpfC9"
   },
   "source": [
    "# 0. Setup Paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "146BB11JpfDA"
   },
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "42hJEdo_pfDB"
   },
   "outputs": [],
   "source": [
    "CUSTOM_MODEL_NAME = 'my_ssd_mobnet' \n",
    "PRETRAINED_MODEL_NAME = 'ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8'\n",
    "PRETRAINED_MODEL_URL = 'http://download.tensorflow.org/models/object_detection/tf2/20200711/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.tar.gz'\n",
    "TF_RECORD_SCRIPT_NAME = 'generate_tfrecord.py'\n",
    "LABEL_MAP_NAME = 'label_map.pbtxt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hbPhYVy_pfDB"
   },
   "outputs": [],
   "source": [
    "paths = {\n",
    "    'WORKSPACE_PATH': os.path.join('Tensorflow', 'workspace'),\n",
    "    'SCRIPTS_PATH': os.path.join('Tensorflow','scripts'),\n",
    "    'APIMODEL_PATH': os.path.join('Tensorflow','models'),\n",
    "    'ANNOTATION_PATH': os.path.join('Tensorflow', 'workspace','annotations'),\n",
    "    'IMAGE_PATH': os.path.join('Tensorflow', 'workspace','images'),\n",
    "    'MODEL_PATH': os.path.join('Tensorflow', 'workspace','models'),\n",
    "    'PRETRAINED_MODEL_PATH': os.path.join('Tensorflow', 'workspace','pre-trained-models'),\n",
    "    'CHECKPOINT_PATH': os.path.join('Tensorflow', 'workspace','models',CUSTOM_MODEL_NAME), \n",
    "    'OUTPUT_PATH': os.path.join('Tensorflow', 'workspace','models',CUSTOM_MODEL_NAME, 'export'), \n",
    "    'TFJS_PATH':os.path.join('Tensorflow', 'workspace','models',CUSTOM_MODEL_NAME, 'tfjsexport'), \n",
    "    'TFLITE_PATH':os.path.join('Tensorflow', 'workspace','models',CUSTOM_MODEL_NAME, 'tfliteexport'), \n",
    "    'PROTOC_PATH':os.path.join('Tensorflow','protoc')\n",
    " }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LwhWZMI0pfDC"
   },
   "outputs": [],
   "source": [
    "files = {\n",
    "    'PIPELINE_CONFIG':os.path.join('Tensorflow', 'workspace','models', CUSTOM_MODEL_NAME, 'pipeline.config'),\n",
    "    'TF_RECORD_SCRIPT': os.path.join(paths['SCRIPTS_PATH'], TF_RECORD_SCRIPT_NAME), \n",
    "    'LABELMAP': os.path.join(paths['ANNOTATION_PATH'], LABEL_MAP_NAME)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HR-TfDGrpfDC"
   },
   "outputs": [],
   "source": [
    "for path in paths.values():\n",
    "    if not os.path.exists(path):\n",
    "        if os.name == 'posix':\n",
    "            !mkdir -p {path}\n",
    "        if os.name == 'nt':\n",
    "            !mkdir {path}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OLU-rs_ipfDE"
   },
   "source": [
    "# 1. Download TF Models Pretrained Models from Tensorflow Model Zoo and Install TFOD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://www.tensorflow.org/install/source_windows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "K-Cmz2edpfDE",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if os.name=='nt':\n",
    "    !pip install wget\n",
    "    import wget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iA1DIq5OpfDE"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(os.path.join(paths['APIMODEL_PATH'], 'research', 'object_detection')):\n",
    "    !git clone https://github.com/tensorflow/models {paths['APIMODEL_PATH']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rJjMHbnDs3Tv"
   },
   "outputs": [],
   "source": [
    "# Install Tensorflow Object Detection \n",
    "if os.name=='posix':  \n",
    "    !apt-get install protobuf-compiler\n",
    "    !cd Tensorflow/models/research && protoc object_detection/protos/*.proto --python_out=. && cp object_detection/packages/tf2/setup.py . && python -m pip install . \n",
    "    \n",
    "if os.name=='nt':\n",
    "    url=\"https://github.com/protocolbuffers/protobuf/releases/download/v3.15.6/protoc-3.15.6-win64.zip\"\n",
    "    wget.download(url)\n",
    "    !move protoc-3.15.6-win64.zip {paths['PROTOC_PATH']}\n",
    "    !cd {paths['PROTOC_PATH']} && tar -xf protoc-3.15.6-win64.zip\n",
    "    os.environ['PATH'] += os.pathsep + os.path.abspath(os.path.join(paths['PROTOC_PATH'], 'bin'))   \n",
    "    !cd Tensorflow/models/research && protoc object_detection/protos/*.proto --python_out=. && copy object_detection\\\\packages\\\\tf2\\\\setup.py setup.py && python setup.py build && python setup.py install\n",
    "    !cd Tensorflow/models/research/slim && pip install -e . "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "VERIFICATION_SCRIPT = os.path.join(paths['APIMODEL_PATH'], 'research', 'object_detection', 'builders', 'model_builder_tf2_test.py')\n",
    "# Verify Installation\n",
    "!python {VERIFICATION_SCRIPT}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tensorflow==2.4.1 tensorflow-gpu==2.4.1 --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip uninstall protobuf matplotlib -y\n",
    "!pip install protobuf matplotlib==3.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install Pillow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pyyaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import object_detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "csofht2npfDE",
    "outputId": "ff5471b2-bed2-43f2-959c-327a706527b6"
   },
   "outputs": [],
   "source": [
    "if os.name =='posix':\n",
    "    !wget {PRETRAINED_MODEL_URL}\n",
    "    !mv {PRETRAINED_MODEL_NAME+'.tar.gz'} {paths['PRETRAINED_MODEL_PATH']}\n",
    "    !cd {paths['PRETRAINED_MODEL_PATH']} && tar -zxvf {PRETRAINED_MODEL_NAME+'.tar.gz'}\n",
    "if os.name == 'nt':\n",
    "    wget.download(PRETRAINED_MODEL_URL)\n",
    "    !move {PRETRAINED_MODEL_NAME+'.tar.gz'} {paths['PRETRAINED_MODEL_PATH']}\n",
    "    !cd {paths['PRETRAINED_MODEL_PATH']} && tar -zxvf {PRETRAINED_MODEL_NAME+'.tar.gz'}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M5KJTnkfpfDC"
   },
   "source": [
    "# 2. Create Label Map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "p1BVDWo7pfDC"
   },
   "outputs": [],
   "source": [
    "labels = [{'name':'licence', 'id':1}]\n",
    "\n",
    "with open(files['LABELMAP'], 'w') as f:\n",
    "    for label in labels:\n",
    "        f.write('item { \\n')\n",
    "        f.write('\\tname:\\'{}\\'\\n'.format(label['name']))\n",
    "        f.write('\\tid:{}\\n'.format(label['id']))\n",
    "        f.write('}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C88zyVELpfDC"
   },
   "source": [
    "# 3. Create TF records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kvf5WccwrFGq",
    "outputId": "49902aeb-0bd7-4298-e1a0-5b4a64eb2064"
   },
   "outputs": [],
   "source": [
    "# OPTIONAL IF RUNNING ON COLAB\n",
    "ARCHIVE_FILES = os.path.join(paths['IMAGE_PATH'], 'archive.tar.gz')\n",
    "if os.path.exists(ARCHIVE_FILES):\n",
    "  !tar -zxvf {ARCHIVE_FILES}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KWpb_BVUpfDD",
    "outputId": "56ce2a3f-3933-4ee6-8a9d-d5ec65f7d73c"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(files['TF_RECORD_SCRIPT']):\n",
    "    !git clone https://github.com/nicknochnack/GenerateTFRecord {paths['SCRIPTS_PATH']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pytz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UPFToGZqpfDD",
    "outputId": "0ebb456f-aadc-4a1f-96e6-fbfec1923e1c"
   },
   "outputs": [],
   "source": [
    "!python {files['TF_RECORD_SCRIPT']} -x {os.path.join(paths['IMAGE_PATH'], 'train')} -l {files['LABELMAP']} -o {os.path.join(paths['ANNOTATION_PATH'], 'train.record')} \n",
    "!python {files['TF_RECORD_SCRIPT']} -x {os.path.join(paths['IMAGE_PATH'], 'test')} -l {files['LABELMAP']} -o {os.path.join(paths['ANNOTATION_PATH'], 'test.record')} "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qT4QU7pLpfDE"
   },
   "source": [
    "# 4. Copy Model Config to Training Folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cOjuTFbwpfDF"
   },
   "outputs": [],
   "source": [
    "if os.name =='posix':\n",
    "    !cp {os.path.join(paths['PRETRAINED_MODEL_PATH'], PRETRAINED_MODEL_NAME, 'pipeline.config')} {os.path.join(paths['CHECKPOINT_PATH'])}\n",
    "if os.name == 'nt':\n",
    "    !copy {os.path.join(paths['PRETRAINED_MODEL_PATH'], PRETRAINED_MODEL_NAME, 'pipeline.config')} {os.path.join(paths['CHECKPOINT_PATH'])}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ga8gpNslpfDF"
   },
   "source": [
    "# 5. Update Config For Transfer Learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z9hRrO_ppfDF"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from object_detection.utils import config_util\n",
    "from object_detection.protos import pipeline_pb2\n",
    "from google.protobuf import text_format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "c2A0mn4ipfDF"
   },
   "outputs": [],
   "source": [
    "config = config_util.get_configs_from_pipeline_file(files['PIPELINE_CONFIG'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uQA13-afpfDF",
    "outputId": "907496a4-a39d-4b13-8c2c-e5978ecb1f10"
   },
   "outputs": [],
   "source": [
    "config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9vK5lotDpfDF"
   },
   "outputs": [],
   "source": [
    "pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()\n",
    "with tf.io.gfile.GFile(files['PIPELINE_CONFIG'], \"r\") as f:                                                                                                                                                                                                                     \n",
    "    proto_str = f.read()                                                                                                                                                                                                                                          \n",
    "    text_format.Merge(proto_str, pipeline_config)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rP43Ph0JpfDG"
   },
   "outputs": [],
   "source": [
    "pipeline_config.model.ssd.num_classes = len(labels)\n",
    "pipeline_config.train_config.batch_size = 4\n",
    "pipeline_config.train_config.fine_tune_checkpoint = os.path.join(paths['PRETRAINED_MODEL_PATH'], PRETRAINED_MODEL_NAME, 'checkpoint', 'ckpt-0')\n",
    "pipeline_config.train_config.fine_tune_checkpoint_type = \"detection\"\n",
    "pipeline_config.train_input_reader.label_map_path= files['LABELMAP']\n",
    "pipeline_config.train_input_reader.tf_record_input_reader.input_path[:] = [os.path.join(paths['ANNOTATION_PATH'], 'train.record')]\n",
    "pipeline_config.eval_input_reader[0].label_map_path = files['LABELMAP']\n",
    "pipeline_config.eval_input_reader[0].tf_record_input_reader.input_path[:] = [os.path.join(paths['ANNOTATION_PATH'], 'test.record')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oJvfgwWqpfDG"
   },
   "outputs": [],
   "source": [
    "config_text = text_format.MessageToString(pipeline_config)                                                                                                                                                                                                        \n",
    "with tf.io.gfile.GFile(files['PIPELINE_CONFIG'], \"wb\") as f:                                                                                                                                                                                                                     \n",
    "    f.write(config_text)   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Zr3ON7xMpfDG"
   },
   "source": [
    "# 6. Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "B-Y2UQmQpfDG"
   },
   "outputs": [],
   "source": [
    "TRAINING_SCRIPT = os.path.join(paths['APIMODEL_PATH'], 'research', 'object_detection', 'model_main_tf2.py')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jMP2XDfQpfDH"
   },
   "outputs": [],
   "source": [
    "command = \"python {} --model_dir={} --pipeline_config_path={} --num_train_steps=10000\".format(TRAINING_SCRIPT, paths['CHECKPOINT_PATH'],files['PIPELINE_CONFIG'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "A4OXXi-ApfDH",
    "outputId": "117a0e83-012b-466e-b7a6-ccaa349ac5ab"
   },
   "outputs": [],
   "source": [
    "print(command)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "i3ZsJR-qpfDH",
    "outputId": "cabec5e1-45e6-4f2f-d9cf-297d9c1d0225"
   },
   "outputs": [],
   "source": [
    "!{command}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4_YRZu7npfDH"
   },
   "source": [
    "# 7. Evaluate the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "80L7-fdPpfDH"
   },
   "outputs": [],
   "source": [
    "command = \"python {} --model_dir={} --pipeline_config_path={} --checkpoint_dir={}\".format(TRAINING_SCRIPT, paths['CHECKPOINT_PATH'],files['PIPELINE_CONFIG'], paths['CHECKPOINT_PATH'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "lYsgEPx9pfDH",
    "outputId": "8632d48b-91d2-45d9-bcb8-c1b172bf6eed"
   },
   "outputs": [],
   "source": [
    "print(command)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lqTV2jGBpfDH"
   },
   "outputs": [],
   "source": [
    "!{command}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "orvRk02UpfDI"
   },
   "source": [
    "# 8. Load Train Model From Checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8TYk4_oIpfDI"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "from object_detection.utils import label_map_util\n",
    "from object_detection.utils import visualization_utils as viz_utils\n",
    "from object_detection.builders import model_builder\n",
    "from object_detection.utils import config_util"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prevent GPU complete consumption\n",
    "gpus = tf.config.list_physical_devices('GPU')\n",
    "if gpus:\n",
    "    try: \n",
    "        tf.config.experimental.set_virtual_device_configuration(\n",
    "            gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=5120)])\n",
    "    except RunTimeError as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tDnQg-cYpfDI"
   },
   "outputs": [],
   "source": [
    "# Load pipeline config and build a detection model\n",
    "configs = config_util.get_configs_from_pipeline_file(files['PIPELINE_CONFIG'])\n",
    "detection_model = model_builder.build(model_config=configs['model'], is_training=False)\n",
    "\n",
    "# Restore checkpoint\n",
    "ckpt = tf.compat.v2.train.Checkpoint(model=detection_model)\n",
    "ckpt.restore(os.path.join(paths['CHECKPOINT_PATH'], 'ckpt-11')).expect_partial()\n",
    "\n",
    "@tf.function\n",
    "def detect_fn(image):\n",
    "    image, shapes = detection_model.preprocess(image)\n",
    "    prediction_dict = detection_model.predict(image, shapes)\n",
    "    detections = detection_model.postprocess(prediction_dict, shapes)\n",
    "    return detections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0EmsmbBZpfDI"
   },
   "source": [
    "# 9. Detect from an Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Y_MKiuZ4pfDI"
   },
   "outputs": [],
   "source": [
    "import cv2 \n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cBDbIhNapfDI"
   },
   "outputs": [],
   "source": [
    "category_index = label_map_util.create_category_index_from_labelmap(files['LABELMAP'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Lx3crOhOzITB"
   },
   "outputs": [],
   "source": [
    "IMAGE_PATH = os.path.join(paths['IMAGE_PATH'], 'testimg2.jpg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 269
    },
    "id": "Tpzn1SMry1yK",
    "outputId": "c392a2c5-10fe-4fc4-9998-a1d4c7db2bd3"
   },
   "outputs": [],
   "source": [
    "img = cv2.imread(IMAGE_PATH)\n",
    "image_np = np.array(img)\n",
    "\n",
    "input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)\n",
    "detections = detect_fn(input_tensor)\n",
    "\n",
    "num_detections = int(detections.pop('num_detections'))\n",
    "detections = {key: value[0, :num_detections].numpy()\n",
    "              for key, value in detections.items()}\n",
    "detections['num_detections'] = num_detections\n",
    "\n",
    "# detection_classes should be ints.\n",
    "detections['detection_classes'] = detections['detection_classes'].astype(np.int64)\n",
    "\n",
    "label_id_offset = 1\n",
    "image_np_with_detections = image_np.copy()\n",
    "\n",
    "viz_utils.visualize_boxes_and_labels_on_image_array(\n",
    "            image_np_with_detections,\n",
    "            detections['detection_boxes'],\n",
    "            detections['detection_classes']+label_id_offset,\n",
    "            detections['detection_scores'],\n",
    "            category_index,\n",
    "            use_normalized_coordinates=True,\n",
    "            max_boxes_to_draw=5,\n",
    "            min_score_thresh=.8,\n",
    "            agnostic_mode=False)\n",
    "\n",
    "plt.imshow(cv2.cvtColor(image_np_with_detections, cv2.COLOR_BGR2RGB))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "detections.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Apply OCR to Detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install easyocr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio===0.8.1 -f https://download.pytorch.org/whl/torch_stable.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import easyocr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "detection_threshold = 0.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = image_np_with_detections\n",
    "scores = list(filter(lambda x: x> detection_threshold, detections['detection_scores']))\n",
    "boxes = detections['detection_boxes'][:len(scores)]\n",
    "classes = detections['detection_classes'][:len(scores)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "width = image.shape[1]\n",
    "height = image.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply ROI filtering and OCR\n",
    "for idx, box in enumerate(boxes):\n",
    "    print(box)\n",
    "    roi = box*[height, width, height, width]\n",
    "    print(roi)\n",
    "    region = image[int(roi[0]):int(roi[2]),int(roi[1]):int(roi[3])]\n",
    "    reader = easyocr.Reader(['en'])\n",
    "    ocr_result = reader.readtext(region)\n",
    "    print(ocr_result)\n",
    "    plt.imshow(cv2.cvtColor(region, cv2.COLOR_BGR2RGB))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for result in ocr_result:\n",
    "    print(np.sum(np.subtract(result[0][2],result[0][1])))\n",
    "    print(result[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OCR Filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "region_threshold = 0.05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_text(region, ocr_result, region_threshold):\n",
    "    rectangle_size = region.shape[0]*region.shape[1]\n",
    "    \n",
    "    plate = [] \n",
    "    for result in ocr_result:\n",
    "        length = np.sum(np.subtract(result[0][1], result[0][0]))\n",
    "        height = np.sum(np.subtract(result[0][2], result[0][1]))\n",
    "        \n",
    "        if length*height / rectangle_size > region_threshold:\n",
    "            plate.append(result[1])\n",
    "    return plate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filter_text(region, ocr_result, region_threshold)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bring it Together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "region_threshold = 0.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ocr_it(image, detections, detection_threshold, region_threshold):\n",
    "    \n",
    "    # Scores, boxes and classes above threhold\n",
    "    scores = list(filter(lambda x: x> detection_threshold, detections['detection_scores']))\n",
    "    boxes = detections['detection_boxes'][:len(scores)]\n",
    "    classes = detections['detection_classes'][:len(scores)]\n",
    "    \n",
    "    # Full image dimensions\n",
    "    width = image.shape[1]\n",
    "    height = image.shape[0]\n",
    "    \n",
    "    # Apply ROI filtering and OCR\n",
    "    for idx, box in enumerate(boxes):\n",
    "        roi = box*[height, width, height, width]\n",
    "        region = image[int(roi[0]):int(roi[2]),int(roi[1]):int(roi[3])]\n",
    "        reader = easyocr.Reader(['en'])\n",
    "        ocr_result = reader.readtext(region)\n",
    "        \n",
    "        text = filter_text(region, ocr_result, region_threshold)\n",
    "        \n",
    "        plt.imshow(cv2.cvtColor(region, cv2.COLOR_BGR2RGB))\n",
    "        plt.show()\n",
    "        print(text)\n",
    "        return text, region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "text, region = ocr_it(image_np_with_detections, detections, detection_threshold, region_threshold)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import uuid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'{}.jpg'.format(uuid.uuid1())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_results(text, region, csv_filename, folder_path):\n",
    "    img_name = '{}.jpg'.format(uuid.uuid1())\n",
    "    \n",
    "    cv2.imwrite(os.path.join(folder_path, img_name), region)\n",
    "    \n",
    "    with open(csv_filename, mode='a', newline='') as f:\n",
    "        csv_writer = csv.writer(f, delimiter=',', quotechar='\"', quoting=csv.QUOTE_MINIMAL)\n",
    "        csv_writer.writerow([img_name, text])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_results(text, region, 'detection_results.csv', 'Detection_Images')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IsNAaYAo0WVL"
   },
   "source": [
    "# 10. Real Time Detections from your Webcam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip uninstall opencv-python-headless -y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o_grs6OGpfDJ"
   },
   "outputs": [],
   "source": [
    "cap = cv2.VideoCapture(0)\n",
    "width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "\n",
    "while cap.isOpened(): \n",
    "    ret, frame = cap.read()\n",
    "    image_np = np.array(frame)\n",
    "    \n",
    "    input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)\n",
    "    detections = detect_fn(input_tensor)\n",
    "    \n",
    "    num_detections = int(detections.pop('num_detections'))\n",
    "    detections = {key: value[0, :num_detections].numpy()\n",
    "                  for key, value in detections.items()}\n",
    "    detections['num_detections'] = num_detections\n",
    "\n",
    "    # detection_classes should be ints.\n",
    "    detections['detection_classes'] = detections['detection_classes'].astype(np.int64)\n",
    "\n",
    "    label_id_offset = 1\n",
    "    image_np_with_detections = image_np.copy()\n",
    "\n",
    "    viz_utils.visualize_boxes_and_labels_on_image_array(\n",
    "                image_np_with_detections,\n",
    "                detections['detection_boxes'],\n",
    "                detections['detection_classes']+label_id_offset,\n",
    "                detections['detection_scores'],\n",
    "                category_index,\n",
    "                use_normalized_coordinates=True,\n",
    "                max_boxes_to_draw=5,\n",
    "                min_score_thresh=.8,\n",
    "                agnostic_mode=False)\n",
    "    \n",
    "    try: \n",
    "        text, region = ocr_it(image_np_with_detections, detections, detection_threshold, region_threshold)\n",
    "        save_results(text, region, 'realtimeresults.csv', 'Detection_Images')\n",
    "    except:\n",
    "        pass\n",
    "\n",
    "    cv2.imshow('object detection',  cv2.resize(image_np_with_detections, (800, 600)))\n",
    "    \n",
    "    if cv2.waitKey(10) & 0xFF == ord('q'):\n",
    "        cap.release()\n",
    "        cv2.destroyAllWindows()\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rzlM4jt0pfDJ"
   },
   "source": [
    "# 10. Freezing the Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "n4olHB2npfDJ"
   },
   "outputs": [],
   "source": [
    "FREEZE_SCRIPT = os.path.join(paths['APIMODEL_PATH'], 'research', 'object_detection', 'exporter_main_v2.py ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0AjO93QDpfDJ"
   },
   "outputs": [],
   "source": [
    "command = \"python {} --input_type=image_tensor --pipeline_config_path={} --trained_checkpoint_dir={} --output_directory={}\".format(FREEZE_SCRIPT ,files['PIPELINE_CONFIG'], paths['CHECKPOINT_PATH'], paths['OUTPUT_PATH'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "F6Lsp3tCpfDJ",
    "outputId": "c3828529-bf06-4df5-d7f3-145890ec3edd"
   },
   "outputs": [],
   "source": [
    "print(command)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1Sw1ULgHpfDJ",
    "outputId": "6fd441e1-9fc9-4889-d072-3395c21e40b6"
   },
   "outputs": [],
   "source": [
    "!{command}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wTPmdqaXpfDK"
   },
   "source": [
    "# 11. Conversion to TFJS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "gZ6UzY_fpfDK",
    "outputId": "0c84722e-1c2b-4002-d857-80827ade828a",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install tensorflowjs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0oxbVynHpfDK"
   },
   "outputs": [],
   "source": [
    "command = \"tensorflowjs_converter --input_format=tf_saved_model --output_node_names='detection_boxes,detection_classes,detection_features,detection_multiclass_scores,detection_scores,num_detections,raw_detection_boxes,raw_detection_scores' --output_format=tfjs_graph_model --signature_name=serving_default {} {}\".format(os.path.join(paths['OUTPUT_PATH'], 'saved_model'), paths['TFJS_PATH'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DB2AGNmJpfDK",
    "outputId": "fbc9f747-f511-47e8-df8f-5ea65cef0374"
   },
   "outputs": [],
   "source": [
    "print(command)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "K7rfT4-hpfDK",
    "outputId": "532707fd-6feb-4bc6-84a3-325b5d16303c"
   },
   "outputs": [],
   "source": [
    "!{command}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o8_hm-itpfDK"
   },
   "outputs": [],
   "source": [
    "# Test Code: https://github.com/nicknochnack/RealTimeSignLanguageDetectionwithTFJS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VtUw73FHpfDK"
   },
   "source": [
    "# 12. Conversion to TFLite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XviMtewLpfDK"
   },
   "outputs": [],
   "source": [
    "TFLITE_SCRIPT = os.path.join(paths['APIMODEL_PATH'], 'research', 'object_detection', 'export_tflite_graph_tf2.py ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "us86cjC4pfDL"
   },
   "outputs": [],
   "source": [
    "command = \"python {} --pipeline_config_path={} --trained_checkpoint_dir={} --output_directory={}\".format(TFLITE_SCRIPT ,files['PIPELINE_CONFIG'], paths['CHECKPOINT_PATH'], paths['TFLITE_PATH'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "n1r5YO3rpfDL",
    "outputId": "5fcdf7a4-eee2-4365-f1ca-1751968379ea"
   },
   "outputs": [],
   "source": [
    "print(command)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "I-xWpHN8pfDL",
    "outputId": "7f6bacd8-d077-43b5-c131-5b081fba24a4"
   },
   "outputs": [],
   "source": [
    "!{command}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iJfYMbN6pfDL"
   },
   "outputs": [],
   "source": [
    "FROZEN_TFLITE_PATH = os.path.join(paths['TFLITE_PATH'], 'saved_model')\n",
    "TFLITE_MODEL = os.path.join(paths['TFLITE_PATH'], 'saved_model', 'detect.tflite')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "command = \"tflite_convert \\\n",
    "--saved_model_dir={} \\\n",
    "--output_file={} \\\n",
    "--input_shapes=1,300,300,3 \\\n",
    "--input_arrays=normalized_input_image_tensor \\\n",
    "--output_arrays='TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3' \\\n",
    "--inference_type=FLOAT \\\n",
    "--allow_custom_ops\".format(FROZEN_TFLITE_PATH, TFLITE_MODEL, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "E8GwUeoFpfDL",
    "outputId": "fac43ea4-cc85-471b-a362-e994b06fd583"
   },
   "outputs": [],
   "source": [
    "print(command)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Nbd7gqHMpfDL",
    "outputId": "7c8fe6d5-2415-4641-8548-39d425c202f7"
   },
   "outputs": [],
   "source": [
    "!{command}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5NQqZRdA21Uc"
   },
   "source": [
    "# 13. Zip and Export Models "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tTVTGCQp2ZJJ"
   },
   "outputs": [],
   "source": [
    "!tar -czf models.tar.gz {paths['CHECKPOINT_PATH']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "whShhB0x3PYJ",
    "outputId": "b773201d-35c9-46a8-b893-4a76bd4d5d97"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "3. Training and Detection.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "anprsys",
   "language": "python",
   "name": "anprsys"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
